# -*- coding: utf-8 -*-
'''
Created on 06.07.2020

@author: yu03
'''
import imageio
import numpy as np
from scipy import signal
import matplotlib.pyplot as plt
import glob
import os
import shutil
from matplotlib.gridspec import GridSpec
import matplotlib.patches as patches
from plt_style import *
from astropy._erfa.core import DPI
colors = color_set_medium

j = complex(0, 1)
c = 3e8 # 光速 [m/s]
Lamda = 633e-9 # 光波长 [m]
Fc = c / Lamda # 光频率 [Hz]

pix_size = 5.3e-6
pix_num = 1280
screen_diameter = pix_num * pix_size

w_0 = 1.12e-3
Z_R = np.pi * w_0**2 / Lamda

dx = np.linspace(0, (pix_num-1)*pix_size, num=pix_num)
dy = dx
V_r_x, V_r_y, V_r_z = 0, 0, 1 # Reference vector
L, D = 0.1, 0.4  # D = displacement, L is static
V_m_x, V_m_y, V_z_m = 0, 0.000, 1 # Measurement vector

M, N, L = 0.0165, 0.033, 0.12

X, Y = np.meshgrid(dx, dy)

def fit_func(x, a, b):
    return a*x+b
 
def w_z(w_0, Z_R, Z):
    return w_0*np.sqrt(1 + (Z/Z_R)**2)
 
def R_z(Z_R, Z):
    return Z*(1+(Z_R/Z)**2)
 
def Phi_Gouy(Z_R, Z):
    return np.arctan(Z/Z_R)    
 
def f(X, Y, V_r_x, V_r_y, V_m_x, V_m_y, M=0.0165, N=0.033, L=0.12, D=0.4, w_0=w_0, Z_R=Z_R):
      
    k = 2 * np.pi / Lamda
    I_0 = 512*1.15
      
    Z_p_r = M + N + L * (2/(1-V_r_x**2-V_r_y**2) - 1)
    Z_p_m = M + N + (L+D) * (2/(1-V_m_x**2-V_m_y**2) - 1)
      
    diff_Z_p = D * 2 / (1-V_m_x**2-V_m_y**2) + L * 2 * ((V_m_x**2+V_m_y**2)-(V_r_x**2+V_r_y**2)) / (1-V_m_x**2-V_m_y**2) / (1-V_r_x**2-V_r_y**2)
      
    d_mea = (X*2*V_m_x+Y*2*V_m_y)/(1+V_m_x**2+V_m_y**2) - (L+D)*4*(V_m_x**2+V_m_y**2)/(1-(V_m_x**2+V_m_y**2)**2)
    d_ref = (X*2*V_r_x+Y*2*V_r_y)/(1+V_r_x**2+V_r_y**2) - L*4*(V_r_x**2+V_r_y**2)/(1-(V_r_x**2+V_r_y**2)**2)
      
    r_mea = np.sqrt( X**2 + Y**2 + (L+D)**2 - ((2*X*V_m_x+2*Y*V_m_y+(L+D)*(1-V_m_x**2-V_m_y**2))/(1+V_m_x**2+V_m_y**2))**2 )
    r_ref = np.sqrt( X**2 + Y**2 + (L)**2 - ((2*X*V_r_x+2*Y*V_r_y+L*(1-V_r_x**2-V_r_y**2))/(1+V_r_x**2+V_r_y**2))**2 )
      
    R_ref = R_z(Z_R, Z_p_r+d_ref)
    R_mea = R_z(Z_R, Z_p_m+d_mea)
      
    Phi_gouy_ref = Phi_Gouy(Z_R, Z_p_r+d_ref)
    Phi_gouy_mea = Phi_Gouy(Z_R, Z_p_m+d_mea)
      
    diff_phi = k * (diff_Z_p + (d_mea-d_ref) + r_mea**2/2/R_mea - r_ref**2/2/R_ref) + Phi_gouy_mea - Phi_gouy_ref
    #     A_beat = 0.5 * I_0 * w_0**2 / w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea) * np.exp((-X**2-Y**2)/w_z(w_0, Z_R, Z_p_r+d_ref) / w_z(w_0, Z_R, Z_p_m+d_mea))
    I_beat = 1/w_z(w_0, Z_R, Z_p_r+d_ref)**2*np.exp(-2*r_ref**2/w_z(w_0, Z_R, Z_p_r+d_ref)**2) + 1/w_z(w_0, Z_R, Z_p_m+d_mea)**2*np.exp(-2*r_mea**2/w_z(w_0, Z_R, Z_p_m+d_mea)**2) + 2/w_z(w_0, Z_R, Z_p_r+d_ref)/w_z(w_0, Z_R, Z_p_m+d_mea)*np.exp(-r_ref**2/w_z(w_0, Z_R, Z_p_r+d_ref)**2-r_mea**2/w_z(w_0, Z_R, Z_p_m+d_mea)**2)*np.cos(diff_phi)
    I_beat = 0.5*I_0*w_0**2*I_beat
#     I_beat = I_beat.round().astype(int)
    return I_beat

 
def center_S(V_r_x=0, V_r_y=0, V_m_x=0, V_m_y=0, L=0.12,D=0):
    center_ref = [2*L*V_r_x/(1-V_r_x**2-V_r_y**2),2*L*V_r_y/(1-V_r_x**2-V_r_y**2)]
    center_mea = [2*(L+D)*V_m_x/(1-V_m_x**2-V_m_y**2),2*(L+D)*V_m_y/(1-V_m_x**2-V_m_y**2)]
    center_ref = [i/pix_size+640 for i in center_ref]
    center_mea = [i/pix_size+640 for i in center_mea]
    center = [(center_ref[0]+center_mea[0])/2, (center_ref[1]+center_mea[1])/2]
    hor_index, ver_index = int(round(center[1])), int(round(center[0]))
    return center_ref, center_mea, hor_index, ver_index
 
screen_diameter = pix_num * pix_size
dx = np.linspace(-0.5*(pix_num-1)*pix_size, 0.5*(pix_num-1)*pix_size, num=pix_num)
dy = dx
X, Y = np.meshgrid(dx, dy)

path = r'C:\Users\yu03\Desktop\Angle_range'
images = []
'''
    Generating image set
'''

def layout_plot(num,unit):
    fig = plt.figure(dpi=150)
    fig.set_size_inches(6, 3, True)
    ax1 = fig.add_subplot(1,2,1)
    ax2 = fig.add_subplot(1,2,2)
    image = f(X, Y, V_r_x, V_r_y, V_m_x, V_m_y,M=M, N=N, L=L, D=D, w_0=w_0, Z_R=Z_R)
    img = ax2.imshow(image,cmap='gray',extent=[0,1280,0,1280],origin='lower', animated=True)
    img.set_clim([0,1000])
    
    center_ref, center_mea, hor_index, ver_index = center_S(V_r_x=V_r_x, V_r_y=V_r_y, V_m_x=V_m_x, V_m_y=V_m_y, L=L,D=D)
#     hor_index, ver_index = 805, 805

    ax2.plot([0,pix_num],[hor_index,hor_index],color=colors['grey'],linewidth=0.3)
    ax2.plot([ver_index,ver_index],[0,pix_num],color=colors['grey'],linewidth=0.3)
    line, = ax2.plot((image[hor_index]//3+50), color='white',linewidth=.5)
    line_ver, = ax2.plot(image[:,ver_index]//3+50, np.arange(pix_num), color='white',linewidth=.5)
    ax2.plot(center_ref[0], center_ref[1],color=color_set_normal['blue'],marker='o')
    ax2.plot(center_mea[0], center_mea[1],color=color_set_normal['red'],marker='x',mew=2)
#     print(center_mea[0]-640, center_mea[1]-640)
    ax2.axis('off')
    
    up, down = 329.7*np.sqrt(2), 74.567
    limit = 584.8-down
    up = up + up -down
    ax1.set_xlim(-up*1.1,up*1.1)
    ax1.set_ylim(-up*1.1,up*1.1)
    ax1.set_aspect('equal')
    ax1.spines['bottom'].set_visible(False)
    ax1.spines['left'].set_visible(False)
    ax1.tick_params(axis='both',which='both',labelbottom=[],labelleft=[],length=0)
    ax1.add_patch(patches.Circle((0, 0),up,facecolor=color_set_light['grey'],fill=True))
    ax1.add_patch(patches.Rectangle((-up, -down),up*2,down*2,facecolor=colors['white'],fill=True))
    ax1.add_patch(patches.Rectangle((-down, -up),down*2,up*2,facecolor=colors['white'],fill=True))
    
    ax1.arrow(up/1.1, 0, up/10, 0, length_includes_head=True, width=up/20,head_width=up/15, color=colors['black'])
    ax1.arrow(0, up/1.1, 0, up/10, length_includes_head=True, width=up/20,head_width=up/15,color=colors['black'])
    ax1.plot([-up,-down/1.5],[0,0],linewidth=1,color=colors['black'])
    ax1.plot([down/1.5,up],[0,0],linewidth=1,color=colors['black'])
    ax1.plot([0,0],[-up,-down/1.5],linewidth=1,color=colors['black'])
    ax1.plot([0,0],[down/1.5,up],linewidth=1,color=colors['black'])
# #     ax1.text(0.65, 0.75, text, fontsize=15, horizontalalignment='left', verticalalignment='center', color=color_set_normal['red'], transform=ax1.transAxes)
    range = ax1.add_patch(patches.Rectangle((down, down),limit,limit,edgecolor=colors['black'],facecolor='none',fill=True))
    ref = ax1.add_patch(patches.Circle((center_ref[0]-640, center_ref[1]-640),down/1.5,facecolor=color_set_normal['blue'],fill=True))
    mea = ax1.plot((center_mea[0]-640)*2-down, (center_mea[1]-640)*2-down,color=color_set_normal['red'],marker='x',mew=1.5,ms=6)
    
    ax1.text(0.55, 0.85, '${P_3}$', fontsize=9, horizontalalignment='left', verticalalignment='center', color=color_set_normal['black'], transform=ax1.transAxes)
    ax1.text(0.82, 0.85, '${P_2}$', fontsize=9, horizontalalignment='left', verticalalignment='center', color=color_set_normal['black'], transform=ax1.transAxes)
    ax1.text(0.82, 0.58, '${P_1}$', fontsize=9, horizontalalignment='left', verticalalignment='center', color=color_set_normal['black'], transform=ax1.transAxes)
    ax1.text(0.55, 0.58, '${P_0}$', fontsize=9, horizontalalignment='left', verticalalignment='center', color=color_set_normal['black'], transform=ax1.transAxes)
#     print(up,down,(center_mea[1]-640)*2-down,(up)/np.sqrt(2))
# #     mea = ax1.plot(1000, -100,color=color_set_normal['red'],marker='x',mew=1.5,ms=6)
#     plt.show()
    return
    
V_m_x, V_m_y = 380e-6, 380e-6   
if 1:
    unit=10
    layout_plot(num=0,unit=unit)
    angle = 1.3e-3
    for num in range(int(unit*8)):
        print(num)
        if num < unit:
            V_m_y += angle/unit   
        elif num >= unit*1 and num < unit*2:
            pass
        
        elif num >= unit*2 and num < unit*3:
            V_m_x += angle/unit   
        elif num >= unit*3 and num < unit*4:
            pass    
            
        elif num >= unit*4 and num < unit*5:
            V_m_y += -angle/unit   
        elif num >= unit*5 and num < unit*6:
            pass

        elif num >=unit*6 and num <unit*7:
            V_m_x += -angle/unit   
        elif num >= unit*7 and num < unit*8:
            pass
         
        layout_plot(num=num,unit=unit)
        plt.savefig(path+'\\'+'%s.jpg'%num)
        plt.clf()
        plt.close()

'''
    Reading images
'''   
if 1:
    files = sorted(glob.glob(path+r'\*.jpg'), key=os.path.getmtime)
    for file in files: #对路径下的每个.pdf文件进行操作
#         print(file)
        img = imageio.imread(file)
        images.append(img)
    print(len(images))
    imageio.mimsave(r'C:\Users\yu03\Desktop\angle_range.gif', images, duration=0.1)